import os
import time
import torch
import numpy as np
from PIL import Image
from PIL.PngImagePlugin import PngInfo
import json
import random
import logging
import re # Import re for regex operations
from typing import Tuple, Dict, Any, Optional, List
import sys # For flushing stdout/stderr

# Import ComfyUI types for better type hinting and autocomplete
from comfy.comfy_types import ComfyNodeABC, IO, InputTypeDict

# Assuming ComfyUI's folder_paths and cli_args are available in the node environment
import folder_paths
from comfy.cli_args import args

from .utils import tensor_to_pil, sanitize_filename, ensure_utf8_friendly, SUCCESS_HIGHLIGHT

# Setup logger for this module
logger = logging.getLogger(__name__)

# Helper to ensure print statements appear immediately in the terminal
def debug_print(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)
    sys.stderr.flush()

class SaveImageEnhancedNode(ComfyNodeABC): # Inherit from ComfyNodeABC
    """
    A ComfyUI node for saving images with enhanced options, including
    custom output folders, flexible filename prefixes, and optional caption saving.
    """
    def __init__(self):
        # Default output directory relative to ComfyUI's output
        self.output_dir = folder_paths.get_output_directory()
        self.type = "output"
        self.compress_level = 4 # Default PNG compression level

    @classmethod
    def INPUT_TYPES(s) -> InputTypeDict: # Use InputTypeDict for type hinting
        return {
            "required": {
                "images": (IO.IMAGE,), # Use IO.IMAGE
                "filename_prefix": (IO.STRING, {"default": "ComfyUI_DN_%date:yyyy-MM-dd%", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes. Supports %batch_num%."}),
                "output_folder": (IO.STRING, {"default": "output", "tooltip": "Subfolder within ComfyUI's output directory, or an absolute path."}),
                "add_counter_suffix": (IO.BOOLEAN, {"default": True, "tooltip": "If True, adds an incrementing numerical suffix (_00001_) to the filename to prevent overwriting. If the filename_prefix ends with a number, that number is treated as the initial counter value."}),
                "counter_padding": (IO.INT, {"default": 5, "min": 1, "max": 9, "step": 1, "tooltip": "Controls the minimum number of digits (with leading zeros) for the counter suffix added by 'add_counter_suffix'. For example, padding of 3 will produce _001_, _002_, etc."}),
                "allow_overwrite": (IO.BOOLEAN, {"default": False, "tooltip": "If True, allows overwriting existing files with the same name. If False, automatically increments the filename prefix to find a unique name."}),
            },
            "optional": {
                "caption_file_extension": (IO.STRING, {"default": ".txt", "tooltip": "The extension for the caption file."}),
                "caption": (IO.STRING, {"forceInput": True, "multiline": True, "default": "", "tooltip": "Optional: String to save as a caption file alongside the image."}),
            },
            "hidden": {
                "prompt": "PROMPT", # Used for saving metadata in PNG
                "extra_pnginfo": "EXTRA_PNGINFO" # Used for saving metadata in PNG
            },
        }

    RETURN_TYPES: Tuple[str] = (IO.STRING,) # Use IO.STRING
    RETURN_NAMES: Tuple[str] = ("last_filename_saved",)
    FUNCTION: str = "save_images" # Use type hint for FUNCTION

    OUTPUT_NODE: bool = True # Use type hint for OUTPUT_NODE

    CATEGORY: str = "Divergent Nodes 👽/Image" # New category for Divergent Nodes
    DESCRIPTION: str = "Saves the input images to a specified directory with optional caption and filename counter."

    def save_images(self, images: torch.Tensor, output_folder: str, filename_prefix: str="ComfyUI_DN_%date:yyyy-MM-dd%", add_counter_suffix: bool=True, counter_padding: int=5, allow_overwrite: bool=False, prompt: Optional[Dict[str, Any]]=None, extra_pnginfo: Optional[Dict[str, Any]]=None, caption: Optional[str]=None, caption_file_extension: str=".txt") -> Tuple[str]:
        """
        Saves images with enhanced options.
        """
        debug_print(f"[DEBUG] [SaveImageEnhanced] save_images method called.")
        full_output_folder = self._get_full_output_folder(output_folder)
        debug_print(f"[DEBUG] [SaveImageEnhanced] Calculated full_output_folder: '{full_output_folder}'")
        os.makedirs(full_output_folder, exist_ok=True)
        logger.info(f"📁 [SaveImageEnhanced] Saving images to: {full_output_folder}")

        last_filename = ""
        
        # Use ComfyUI's path helper to get base filename and initial counter
        # We will modify the filename later based on add_counter_suffix
        # The counter returned here is based on existing files matching the prefix pattern
        # Strip leading/trailing whitespace from filename_prefix
        cleaned_filename_prefix = filename_prefix.strip()
        
        # Ensure images is not empty to avoid index error
        if not images.shape[0] > 0:
            logger.error("❌ [SaveImageEnhanced] No images provided to the node. Please connect an image input.")
            return ("ERROR: No images provided.",)

        _, base_filename_without_counter, initial_counter, subfolder, _ = folder_paths.get_save_image_path(
            cleaned_filename_prefix, full_output_folder, images[0].shape[1], images[0].shape[0]
        )

        # --- Logic for allow_overwrite and filename increment ---
        # When allow_overwrite is False, we want to increment the user-provided
        # filename_prefix itself (e.g., 000 -> 001 -> 002) for each run of the node,
        # and for batches, we want consecutive numbers (000, 001, 002 for a 3-image batch).
        # 
        # We'll handle this by:
        # 1. Determining the base prefix part (before any auto-incremented number).
        # 2. For the first image in a run, find the next available number in the sequence
        #    by calling _increment_filename_if_needed on a "proposed" filename.
        # 3. For subsequent images in the same run, simply increment from that starting number.
        
        # State variables for the new increment logic
        increment_start_number = None
        increment_format_width = None
        increment_prefix_part = None
        # Flag to know if we've done the initial increment check for this run
        has_done_initial_increment_check = False 
        
        # Pre-calculate the base filename prefix for the batch if needed
        # This is the prefix with all dynamic parts filled in, except the final counter.
        base_filename_prefix_for_batch = None
        if not allow_overwrite:
            # Use a dummy PIL image (first one) to get base prefix, assuming 
            # dimensions and prompts are the same for all in a batch for prefix calc.
            if images.shape[0] > 0:
                dummy_pil_image_for_prefix = tensor_to_pil(images[0])
                if dummy_pil_image_for_prefix:
                    base_filename_prefix_for_batch = self._get_base_filename_prefix(
                        filename_prefix, dummy_pil_image_for_prefix, prompt, extra_pnginfo
                    )
                else:
                    logger.error("❌ [SaveImageEnhanced] Could not convert first tensor to PIL for base prefix calculation.")
                    base_filename_prefix_for_batch = "ERROR_BASE_PREFIX"

        # Use the original counter mechanism for other cases
        current_counter = initial_counter

        for batch_number, image in enumerate(images):
            # Convert tensor to PIL Image
            pil_image = tensor_to_pil(image)
            if pil_image is None:
                logger.error(f"❌ [SaveImageEnhanced] Could not convert tensor for batch {batch_number}. Skipping.")
                continue

            # Handle filename generation based on allow_overwrite setting
            if not allow_overwrite:
                # --- New Logic for allow_overwrite=False ---
                
                if base_filename_prefix_for_batch is None:
                    # Fallback, should not happen if the check above passed
                    base_prefix = "ERROR_FALLBACK"
                else:
                    base_prefix = base_filename_prefix_for_batch
                
                # If this is the very first image being processed in this node execution
                if not has_done_initial_increment_check:
                    # 1. Propose a base filename for the first image (usually the user's prefix)
                    #    We assume the user's prefix might end in a number (like '000').
                    proposed_first_filename_base = base_prefix
                    sanitized_proposed_first = sanitize_filename(proposed_first_filename_base)
                    
                    debug_print(f"[DEBUG] [SaveImageEnhanced] allow_overwrite=False: Base prefix for batch: '{base_prefix}'")
                    debug_print(f"[DEBUG] [SaveImageEnhanced] allow_overwrite=False: Sanitized proposed first filename base: '{sanitized_proposed_first}'")
                    
                    # 2. Use _increment_filename_if_needed to find the actual first unique filename
                    #    This handles the case where '000.png' exists and finds '003.png' or whatever is next.
                    actual_first_filename_with_ext = self._increment_filename_if_needed(
                        sanitized_proposed_first, full_output_folder, ".png"
                    )
                    
                    debug_print(f"[DEBUG] [SaveImageEnhanced] allow_overwrite=False: _increment_filename_if_needed returned: '{actual_first_filename_with_ext}'")
                    
                    # 3. Parse the actual filename returned to get the starting number, format, and prefix
                    actual_first_filename_base, _ = os.path.splitext(actual_first_filename_with_ext)
                    # Regex to find the numeric part at the end
                    match = re.search(r'(\d+)$', actual_first_filename_base)
                    if match:
                        increment_start_number = int(match.group(1))
                        increment_format_width = len(match.group(1))
                        increment_prefix_part = actual_first_filename_base[:match.start()]
                        debug_print(f"[DEBUG] [SaveImageEnhanced] allow_overwrite=False: Parsed increment_start_number: {increment_start_number}, format_width: {increment_format_width}, prefix_part: '{increment_prefix_part}'")
                    else:
                        # If the returned filename doesn't end in a number, 
                        # we treat the whole name as the prefix and start counting from 0 or 1.
                        # Let's append '_0' and try to increment that.
                        logger.info(f"First filename '{actual_first_filename_base}' does not end with a number. Will append '_0' and increment.")
                        increment_start_number = 0
                        increment_format_width = 1 # Default width for appended number
                        increment_prefix_part = f"{actual_first_filename_base}_"
                        # We need to check again for uniqueness with the new base
                        proposed_with_appended = f"{increment_prefix_part}{str(increment_start_number).zfill(increment_format_width)}"
                        final_unique_with_appended = self._increment_filename_if_needed(
                            proposed_with_appended, full_output_folder, ".png"
                        )
                        # Re-parse
                        final_base, _ = os.path.splitext(final_unique_with_appended)
                        match_final = re.search(r'(\d+)$', final_base)
                        if match_final:
                            increment_start_number = int(match_final.group(1))
                            increment_format_width = len(match_final.group(1))
                            increment_prefix_part = final_base[:match_final.start()]
                            debug_print(f"[DEBUG] [SaveImageEnhanced] allow_overwrite=False (after append): Parsed increment_start_number: {increment_start_number}, format_width: {increment_format_width}, prefix_part: '{increment_prefix_part}'")
                        else:
                            # Should not happen if _increment_filename_if_needed works correctly
                            logger.error(f"Could not parse number from final filename '{final_base}' after append. Using '_0' as final name base.")
                            increment_start_number = 0
                            increment_format_width = 1
                            increment_prefix_part = f"{actual_first_filename_base}_"
                            
                    # Mark that we've done the initial check
                    has_done_initial_increment_check = True
                
                # --- For all images (including the first one after initial check) ---
                # 4. Calculate the final filename for the current image based on its batch_number
                #    and the starting number determined above.
                current_image_number = increment_start_number + batch_number
                formatted_current_number = str(current_image_number).zfill(increment_format_width)
                final_filename_base = f"{increment_prefix_part}{formatted_current_number}"
                
                debug_print(f"[DEBUG] [SaveImageEnhanced] allow_overwrite=False: For batch_number {batch_number}, current_image_number: {current_image_number}, formatted: '{formatted_current_number}', final_filename_base: '{final_filename_base}'")
                
                # Sanitize the final base filename
                sanitized_final_base = sanitize_filename(final_filename_base)
                
                # Construct full paths
                full_path_no_ext = os.path.join(full_output_folder, sanitized_final_base)
                full_image_path = f"{full_path_no_ext}.png" # Always save as PNG
                
            else:
                # --- Existing Logic for allow_overwrite=True or default behavior ---
                # Use the original counter mechanism influenced by folder_paths and add_counter_suffix
                # but now with custom counter_padding.
                debug_print(f"[INFO] [SaveImageEnhanced] Using standard ComfyUI counter logic (add_counter_suffix={add_counter_suffix}). If filename_prefix ends with a number, that number is treated as the initial counter value by ComfyUI.")
                
                # Get the base prefix without %batch_num% for consistency and to apply our own counter
                base_prefix_for_counter = self._get_base_filename_prefix(filename_prefix, pil_image, prompt, extra_pnginfo)
                sanitized_base_prefix_for_counter = sanitize_filename(base_prefix_for_counter)
                
                # If add_counter_suffix is True, we need to find the next available number
                # and format it according to counter_padding.
                if add_counter_suffix:
                    debug_print(f"[DEBUG] [SaveImageEnhanced] add_counter_suffix=True, using custom logic with padding {counter_padding}.")
                    # Start with the counter provided by folder_paths.get_save_image_path
                    # This gives us a good starting point based on existing files.
                    temp_counter = current_counter 
                    
                    while True:
                        # Format the counter with the specified padding
                        formatted_counter = str(temp_counter).zfill(counter_padding)
                        # Construct the proposed filename base
                        proposed_filename_base = f"{sanitized_base_prefix_for_counter}_{formatted_counter}"
                        proposed_full_path = os.path.join(full_output_folder, f"{proposed_filename_base}.png")
                        
                        # Check if this file already exists
                        if not os.path.exists(proposed_full_path):
                            # Found a unique name
                            debug_print(f"[DEBUG] [SaveImageEnhanced] Found unique filename: '{proposed_filename_base}.png'")
                            sanitized_filename = proposed_filename_base
                            # We don't increment current_counter here because our loop handles it
                            # and we only want to increment it once per image in the batch.
                            break
                        else:
                            debug_print(f"[DEBUG] [SaveImageEnhanced] Filename '{proposed_filename_base}.png' exists, incrementing counter.")
                            temp_counter += 1
                else:
                    # If add_counter_suffix is False, just use the sanitized base prefix
                    debug_print(f"[DEBUG] [SaveImageEnhanced] add_counter_suffix=False, using base prefix only.")
                    sanitized_filename = sanitized_base_prefix_for_counter
                
                full_path_no_ext = os.path.join(full_output_folder, sanitized_filename)
                full_image_path = f"{full_path_no_ext}.png" # Always save as PNG

                # Increment counter only if suffix is added (original logic)
                # Note: With our custom logic, we only increment current_counter once per image.
                # The internal loop for finding a unique name uses temp_counter.
                if add_counter_suffix:
                    current_counter += 1

            # Save image
            try:
                info = PngInfo() if not args.disable_metadata else None
                if info:
                    metadata_dict = {}
                    if prompt is not None:
                        metadata_dict["prompt"] = prompt
                    if extra_pnginfo is not None:
                        metadata_dict.update(extra_pnginfo)

                    if metadata_dict:
                        # Convert combined metadata to JSON string
                        metadata_json = json.dumps(metadata_dict, indent=2)
                        # Ensure the entire JSON string is UTF-8 friendly
                        info.add_text("prompt", ensure_utf8_friendly(metadata_json))
                
                pil_image.save(full_image_path, pnginfo=info, compress_level=self.compress_level)
                logger.log(SUCCESS_HIGHLIGHT, f"Image saved: {full_image_path}") # Use SUCCESS_HIGHLIGHT
                last_filename = full_image_path
            except Exception as e:
                logger.error(f"❌ [SaveImageEnhanced] Error saving image to '{full_image_path}'. Please check folder permissions or disk space. Error: {e}", exc_info=True)
                last_filename = f"ERROR: Could not save image {full_image_path} - {e}"
                continue # Try to save next image

            # Save caption if provided
            if caption and caption_file_extension:
                caption_filename = f"{full_path_no_ext}{caption_file_extension}"
                try:
                    # Ensure caption is UTF-8 friendly
                    safe_caption = ensure_utf8_friendly(caption)
                    with open(caption_filename, "w", encoding="utf-8") as f:
                        f.write(safe_caption)
                    logger.info(f"📄 [SaveImageEnhanced] Caption saved: {caption_filename}")
                except Exception as e:
                    logger.error(f"❌ [SaveImageEnhanced] Error saving caption to '{caption_filename}'. Please check folder permissions. Error: {e}", exc_info=True)

        return (last_filename,)

    def _get_full_output_folder(self, output_folder_input: str) -> str:
        """
        Determines the full path for the output folder.
        If output_folder_input is an absolute path, use it directly.
        Otherwise, treat it as a subfolder relative to ComfyUI's output directory.
        """
        debug_print(f"[DEBUG] [SaveImageEnhanced] _get_full_output_folder called with: '{output_folder_input}'")
        
        # Handle empty string explicitly to resolve to the base output directory
        # Use self.output_dir directly to ensure it matches exactly what 
        # folder_paths.get_save_image_path expects, including any trailing slashes.
        if output_folder_input == "":
            debug_print(f"[DEBUG] [SaveImageEnhanced] _get_full_output_folder: input was empty, returning self.output_dir: '{self.output_dir}'")
            return self.output_dir
            
        if os.path.isabs(output_folder_input):
            result = output_folder_input
        else:
            result = os.path.join(self.output_dir, output_folder_input)
            
        # Normalize the path to ensure consistent representation (e.g., remove trailing slashes)
        # This might help with the comparison inside folder_paths.get_save_image_path
        result = os.path.normpath(result)
        debug_print(f"[DEBUG] [SaveImageEnhanced] _get_full_output_folder returning: '{result}'")
        return result

    def _get_base_filename_prefix(self, filename_prefix: str, pil_image: Image.Image,
                                  prompt: Optional[Dict[str, Any]], extra_pnginfo: Optional[Dict[str, Any]]) -> str:
        """
        Generates the base part of the filename prefix, performing all dynamic replacements
        except for the batch/counter number (%batch_num%).
        This is used when allow_overwrite is False to determine the base prefix to increment.
        """
        # Replace date placeholders
        filename = filename_prefix.replace("%date:yyyy-MM-dd%", time.strftime("%Y-%m-%d"))
        filename = filename.replace("%date:yyyyMMdd_HHmmss%", time.strftime("%Y%m%d_%H%M%S"))

        # Replace image dimension placeholders
        filename = filename.replace("%width%", str(pil_image.width))
        filename = filename.replace("%height%", str(pil_image.height))

        # NOTE: %batch_num% is NOT replaced here. It will be handled by the caller
        # if needed, or for the new allow_overwrite logic, we will increment the
        # numeric part of the final base prefix.

        # Replace placeholders from other nodes (e.g., %Empty Latent Image.width%)
        if prompt:
            for node_id, node_data in prompt.items():
                if isinstance(node_data, dict) and "inputs" in node_data:
                    for input_name, input_value in node_data["inputs"].items():
                        placeholder = f"%{node_id}.{input_name}%"
                        if placeholder in filename:
                            filename = filename.replace(placeholder, str(input_value))
        
        debug_print(f"[DEBUG] [SaveImageEnhanced] _get_base_filename_prefix: input '{filename_prefix}' -> output '{filename}'")
        return filename

    def _get_image_filename(self, filename_prefix: str, index: int, pil_image: Image.Image,
                            prompt: Optional[Dict[str, Any]], extra_pnginfo: Optional[Dict[str, Any]],
                            add_counter_suffix: bool) -> str:
        """
        Generates a dynamic filename based on the prefix, index, image properties, and prompt info.
        Note: This method is not used in the main save_images loop anymore for the new logic,
        but is kept for potential future use or compatibility.
        """
        # Replace date placeholders
        filename = filename_prefix.replace("%date:yyyy-MM-dd%", time.strftime("%Y-%m-%d"))
        filename = filename.replace("%date:yyyyMMdd_HHmmss%", time.strftime("%Y%m%d_%H%M%S"))

        # Replace image dimension placeholders
        filename = filename.replace("%width%", str(pil_image.width))
        filename = filename.replace("%height%", str(pil_image.height))

        # Replace batch number placeholder
        # Note: %batch_num% is typically 0-indexed for the current batch,
        # while 'index' here is the global counter.
        # If user expects %batch_num% to be 0-indexed for current batch,
        # this needs adjustment. For now, using global 'index'.
        filename = filename.replace("%batch_num%", str(index).zfill(5)) # Pad with zeros

        # Replace placeholders from other nodes (e.g., %Empty Latent Image.width%)
        if prompt:
            for node_id, node_data in prompt.items():
                if isinstance(node_data, dict) and "inputs" in node_data:
                    for input_name, input_value in node_data["inputs"].items():
                        placeholder = f"%{node_id}.{input_name}%"
                        if placeholder in filename:
                            filename = filename.replace(placeholder, str(input_value))
        
        # The counter suffix is now handled by the main loop's `current_counter`
        # and the `add_counter_suffix` logic.
        # This function now just prepares the base filename part.
        return filename

    def _increment_filename_if_needed(self, base_filename: str, full_output_folder: str, extension: str = ".png") -> str:
        """
        Checks if a file with the given base_filename and extension exists in the full_output_folder.
        If it does, and allow_overwrite is False, it attempts to increment a numeric suffix
        in the base_filename to find a unique name.

        Args:
            base_filename (str): The base name of the file (without path or extension).
            full_output_folder (str): The full path to the output directory.
            extension (str): The file extension (e.g., '.png').

        Returns:
            str: A unique filename (base + extension) that does not currently exist in the folder.
        """
        debug_print(f"[DEBUG] [SaveImageEnhanced] _increment_filename_if_needed called with base_filename='{base_filename}', full_output_folder='{full_output_folder}', extension='{extension}'")
        
        candidate_filename = base_filename
        candidate_full_path = os.path.join(full_output_folder, f"{candidate_filename}{extension}")

        # Check if the file already exists
        if os.path.exists(candidate_full_path):
            debug_print(f"[DEBUG] File '{candidate_full_path}' already exists. Attempting to increment filename.")
            
            # Regex to find a trailing numeric part (e.g., '000', '450', '00000')
            # It looks for one or more digits at the end of the string.
            match = re.search(r'(\d+)$', base_filename)
            
            if match:
                # Found a numeric suffix
                numeric_part = match.group(1)
                prefix_part = base_filename[:match.start()]
                # Preserve the original format (e.g., leading zeros)
                format_width = len(numeric_part)
                try:
                    current_number = int(numeric_part)
                except ValueError:
                    # If conversion fails (shouldn't happen with the regex), treat as no number
                    logger.warning(f"Numeric part '{numeric_part}' in filename '{base_filename}' could not be parsed as an integer. Appending '_1'.")
                    candidate_filename = f"{base_filename}_1"
                    candidate_full_path = os.path.join(full_output_folder, f"{candidate_filename}{extension}")
                    while os.path.exists(candidate_full_path):
                        # Simple increment for appended suffix
                        match_simple = re.search(r'_([0-9]+)$', candidate_filename)
                        if match_simple:
                            try:
                                num = int(match_simple.group(1))
                                new_num = num + 1
                                candidate_filename = f"{candidate_filename[:match_simple.start()]}_{new_num}"
                            except ValueError:
                                candidate_filename = f"{candidate_filename}_1" # Fallback
                                break # Prevent infinite loop if parsing fails again
                        else:
                            candidate_filename = f"{candidate_filename}_1"
                            break # Should not happen after first append
                        candidate_full_path = os.path.join(full_output_folder, f"{candidate_filename}{extension}")
                else:
                    # Successfully parsed the number, now increment
                    new_number = current_number + 1
                    # Format the new number with the same width, preserving leading zeros
                    formatted_new_number = str(new_number).zfill(format_width)
                    candidate_filename = f"{prefix_part}{formatted_new_number}"
                    candidate_full_path = os.path.join(full_output_folder, f"{candidate_filename}{extension}")
                    
                    debug_print(f"[DEBUG] [SaveImageEnhanced] _increment_filename_if_needed: Initial increment: '{base_filename}' -> '{candidate_filename}'. Now checking for uniqueness.")
                    
                    # Loop to ensure uniqueness if the incremented name also exists
                    while os.path.exists(candidate_full_path):
                        debug_print(f"[DEBUG] [SaveImageEnhanced] _increment_filename_if_needed: '{candidate_full_path}' also exists. Incrementing further.")
                        # Re-parse the number from the new candidate filename
                        match_retry = re.search(r'(\d+)$', candidate_filename)
                        if match_retry and candidate_filename.startswith(prefix_part):
                            try:
                                current_number = int(match_retry.group(1))
                                new_number = current_number + 1
                                formatted_new_number = str(new_number).zfill(format_width)
                                candidate_filename = f"{prefix_part}{formatted_new_number}"
                                candidate_full_path = os.path.join(full_output_folder, f"{candidate_filename}{extension}")
                                debug_print(f"[DEBUG] [SaveImageEnhanced] _increment_filename_if_needed: Incremented to: '{candidate_filename}'")
                            except ValueError:
                                # If parsing fails during loop, append a simple suffix
                                logger.warning(f"Could not parse number in incremented filename '{candidate_filename}'. Appending '_1'.")
                                candidate_filename = f"{candidate_filename}_1"
                                candidate_full_path = os.path.join(full_output_folder, f"{candidate_filename}{extension}")
                                # Break to avoid potential infinite loop if parsing keeps failing
                                break
                        else:
                            # If the structure is unexpected, append a simple suffix
                            logger.warning(f"Filename structure changed unexpectedly during increment. Appending '_1'.")
                            candidate_filename = f"{candidate_filename}_1"
                            candidate_full_path = os.path.join(full_output_folder, f"{candidate_filename}{extension}")
                            break # Prevent infinite loop if structure is wrong
            
            else:
                # No trailing numeric part found, append '_1' and then increment if needed
                debug_print(f"[DEBUG] No trailing numeric part found in '{base_filename}'. Appending '_1'.")
                candidate_filename = f"{base_filename}_1"
                candidate_full_path = os.path.join(full_output_folder, f"{candidate_filename}{extension}")
                while os.path.exists(candidate_full_path):
                     # Simple increment for appended suffix
                    match_simple = re.search(r'_([0-9]+)$', candidate_filename)
                    if match_simple:
                        try:
                            num = int(match_simple.group(1))
                            new_num = num + 1
                            candidate_filename = f"{candidate_filename[:match_simple.start()]}_{new_num}"
                        except ValueError:
                            candidate_filename = f"{candidate_filename}_1" # Fallback
                            break # Prevent infinite loop if parsing fails again
                    else:
                        candidate_filename = f"{candidate_filename}_1"
                        break # Should not happen after first append
                    candidate_full_path = os.path.join(full_output_folder, f"{candidate_filename}{extension}")

            logger.info(f"Incremented filename to '{candidate_filename}{extension}' to avoid overwrite.")

        # Return the final unique filename (base + extension)
        debug_print(f"[DEBUG] [SaveImageEnhanced] _increment_filename_if_needed returning: '{candidate_filename}{extension}'")
        return f"{candidate_filename}{extension}"


NODE_CLASS_MAPPINGS = {
    "SaveImageEnhancedNode": SaveImageEnhancedNode

}

# Define the master NODE_DISPLAY_NAME_MAPPINGS
NODE_DISPLAY_NAME_MAPPINGS = {
    "SaveImageEnhancedNode": "Save Image Enhanced (DN)"
}